import os.path
import random
from io import BytesIO, UnsupportedOperation, StringIO

from Crypto.Cipher import AES

SECRET_KEY = b"E7lME9qbwfc="
KEY_MAGIC = b"njzf_dsas_240612"
AES_MODE = AES.MODE_CBC
AES_IV: bytes = b"0000000000000000"
SIMPLE_MAGIC_NUM = bytearray(b'\x02\x00\x00\x02\x02\x00\x02\x04')
AES_MAGIC_NUM = bytearray(b'\x02\x00\x00\x02\x02\x00\x02\x05')

class NotReadableError(UnsupportedOperation):
    def __init__(self, error_info="not readable"):
        self.error_info = error_info

    def __str__(self):
        return self.error_info
    
class NotWritableError(UnsupportedOperation):
    def __init__(self, error_info="not writable"):
        self.error_info = error_info

    def __str__(self):
        return self.error_info

class UBytesIO(BytesIO):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.read_enable = True
        self.write_enable = True

    def disable_read(self):
        self.read_enable = False

    def disable_write(self):
        self.write_enable = False
    def read(self, *args, **kwargs):
        if not self.read_enable:
            raise NotReadableError
        return super().read(*args, **kwargs)
    
    def write(self, *args, **kwargs):
        if not self.write_enable:
            raise NotWritableError
        return super().write(*args, **kwargs)


class UStringIO(StringIO):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.read_enable = True
        self.write_enable = True

    def disable_read(self):
        self.read_enable = False

    def disable_write(self):
        self.write_enable = False

    def read(self, *args, **kwargs):
        if not self.read_enable:
            raise NotReadableError
        return super().read(*args, **kwargs)
    
    def readline(self, __size = -1):
        if not self.read_enable:
            raise NotReadableError
        return super().readline(__size)

    def readlines(self, __hint = -1):
        if not self.read_enable:
            raise NotReadableError
        return super().readlines(__hint)

    def write(self, *args, **kwargs):
        if not self.write_enable:
            raise NotWritableError
        return super().write(*args, **kwargs)
    
class FileOpen:
    """
        Open file and return a IO object.  Raise OSError upon failure.

        mode is an optional string that specifies the mode in which the file
        is opened. It defaults to 'r' which means open for reading in text
        mode.  Other common values are 'w' for writing (truncating the file if
        it already exists), 'x' for creating and writing to a new file, and
        'a' for appending (which in this class, means that all writes
        append to the end of the file regardless of the current seek position).
        In text mode, if encoding is not specified, it will be set to UTF-8.
        The available modes are:

        ========= ===============================================================
        Character Meaning
        --------- ---------------------------------------------------------------
        'r'       open for reading (default)
        'w'       open for writing, truncating the file first
        'x'       create a new file and open it for writing
        'a'       open for writing, appending to the end of the file if it exists
        'b'       binary mode (default)
        't'       text mode
        ========= ===============================================================

        The default mode is 'rb' (open for reading binary). The 'x' mode implies
        'w' and raises an `FileExistsError` if the file already exists.

        Python distinguishes between files opened in binary and text modes,
        even when the underlying operating system doesn't. Files opened in
        binary mode (appending 'b' to the mode argument) return contents as
        bytes objects without any decoding. In text mode (the default, or when
        't' is appended to the mode argument), the contents of the file are
        returned as strings, the bytes having been first decoded using a
        default UTF-8 encoding or using the specified encoding if given.

        encoding is the name of the encoding used to decode or encode the
        file. This should only be used in text mode. The default encoding is
        UTF-8.

        Demo:
        if __name__ == "__main__":
            import json
            with FileOpen("/tmp/a.json", "wt") as file:
                json.dump({"我": "张三", "你": "李四"}, file, indent=4, ensure_ascii=False)
            with FileOpen("/tmp/a.json", "rt") as file:
                print(file.read())
            _s = time.time()
            with FileOpen("/tmp/a.txt", "wt", encrypt_mode="aes") as file:
                file.write("张三李四王二麻子\n"*1024*1024*10)
            _m = time.time()
            print(f"write consuming:{_m - _s}")
            with FileOpen("/tmp/a.txt", "rt") as file:
                a = file.read()
                print(len(a)//(1024*1024))
            print(f"read consuming:{time.time() - _m}")
            _s = time.time()
            with FileOpen("2024-09-10.pkl") as file:
                _m = time.time()
                print(f"decrypt consuming:{_m - _s}")
                import pandas as pd
                df = pd.read_pickle(file, compression="zip")
            print(f"read pickle consuming:{time.time() - _m}")
            print(df.info())
        """
    def __init__(self, file_path: str, open_mode: str = "rb", encoding="utf-8", encrypt_mode=""):
        self.encrypt_lt = [7, 8, 3, 5, 6, 9, 1, 4, 2, 0]
        self.decrypt_lt = [9, 6, 8, 2, 7, 3, 4, 0, 1, 5]
        self.ft = 1000
        self.ft_slice = self.ft // 10
        self.encrypt_mode = encrypt_mode
        self.file_path = file_path
        self.open_mode = open_mode[0]
        if len(open_mode) >= 2 and open_mode[1] == 't':
            self.io_type = 't'
            self.io_func = UStringIO
            self.encoding = encoding
        else:
            self.io_type = 'b'
            self.io_func = UBytesIO
            self.encoding = None
        if self.open_mode in ["w", "x", "a"]:
            self.readable = 0
            self.writable = 1
        else:
            self.readable = 1
            self.writable = 0
        self._file_check()
        if self.open_mode in ["a", "r"]:
            with open(file_path, "rb") as f:
                content = f.read()
            content = self.decrypt(content)
            if self.io_type == 't':
                content = content.decode(encoding=self.encoding)
            _seek = self.writable and len(content)
            self.io = self.io_func(content)
            self.io.seek(_seek)
        else:
            self.io = self.io_func()
        if not self.readable:
            self.io.disable_read()
        if not self.writable:
            self.io.disable_write()

    def _file_check(self):
        if os.path.isdir(self.file_path):
            raise IsADirectoryError(f"{self.file_path} is a dir")
        if self.open_mode in ["r", "a"] and not os.path.exists(self.file_path):
            raise FileNotFoundError
        if self.open_mode == "x" and os.path.exists(self.file_path):
            raise FileExistsError


    def encrypt(self, byte_stream: bytes):
        """

        :param byte_stream:
        :return:
        """
        if self.encrypt_mode.lower() == "simple":
            return self.simple_encrypt(byte_stream)
        elif self.encrypt_mode.lower() == "aes":
            return self.aes_encrypt(byte_stream)
        else:
            return byte_stream

    def aes_encrypt(self, byte_stream):
        new_byte_stream = bytearray()
        new_byte_stream = AES_MAGIC_NUM + new_byte_stream
        for _ in range(self.ft):
            new_byte_stream.append(random.randint(1, 255))
        crypto_inst = AES.new(KEY_MAGIC, AES_MODE, AES_IV)
        pad_len = 16 - len(byte_stream) % 16
        data = byte_stream + bytes([pad_len] * pad_len)
        cipher_text = crypto_inst.encrypt(data)
        new_byte_stream = new_byte_stream + cipher_text
        return new_byte_stream

    def simple_encrypt(self, byte_stream):
        """
        对pickle文件进行转换
        0-7
        1-8,
        2-3,
        3-5,
        4-6,
        5-9,
        6-1,
        7-4,
        8-2
        9-0
        :param byte_stream:
        :return:
        """
        new_byte_stream = bytearray()
        for _ in range(self.ft):
            new_byte_stream.append(random.randint(1, 255))
        length = len(byte_stream)
        loop = length // self.ft
        for i in range(loop):
            base_pos = i * self.ft
            for j in range(len(self.encrypt_lt)):
                index = self.encrypt_lt[j]
                new_byte_stream += byte_stream[base_pos + index * self.ft_slice: base_pos + (index + 1) * self.ft_slice]

        new_byte_stream += byte_stream[len(new_byte_stream) - self.ft:]
        new_byte_stream = SIMPLE_MAGIC_NUM + new_byte_stream

        return new_byte_stream

    def decrypt(self, byte_stream:bytes):
        """

        :param byte_stream:
        :return:
        """
        if byte_stream[:len(SIMPLE_MAGIC_NUM)] == SIMPLE_MAGIC_NUM:
            return self.simple_decrypt(byte_stream)
        elif byte_stream[:len(SIMPLE_MAGIC_NUM)] == AES_MAGIC_NUM:
            return self.aes_decrypt(byte_stream)
        else:
            return byte_stream

    def simple_decrypt(self, byte_stream):
        """
        对pickle文件进行转换
        0-7
        1-8,
        2-3,
        3-5,
        4-6,
        5-9,
        6-1,
        7-4,
        8-2
        9-0
        :param byte_stream:
        :return:
        """
        byte_stream = byte_stream[len(SIMPLE_MAGIC_NUM) + self.ft:]
        new_byte_stream = bytearray()
        length = len(byte_stream)
        loop = length // self.ft
        for i in range(loop):
            base_pos = i * self.ft
            for j in range(len(self.decrypt_lt)):
                index = self.decrypt_lt[j]
                new_byte_stream += byte_stream[base_pos + index * self.ft_slice: base_pos + (index + 1) * self.ft_slice]

        new_byte_stream += byte_stream[len(new_byte_stream):]

        return new_byte_stream

    def aes_decrypt(self, byte_stream):
        byte_stream = byte_stream[len(AES_MAGIC_NUM) + self.ft:]
        own_crypto = AES.new(KEY_MAGIC, AES_MODE, AES_IV)
        new_byte_stream = own_crypto.decrypt(byte_stream)
        pad_len = new_byte_stream[-1]
        return new_byte_stream[:-pad_len]


    def close(self):
        if self.writable:
            content = self.io.getvalue()
            if self.io_type == 't':
                content = content.encode(encoding=self.encoding)
            content = self.encrypt(content)
            with open(self.file_path, "wb") as f:
                f.write(content)
        self.io.close()

    def __enter__(self):
        return self.io

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
"""

"""